import numpy as np
import pandas as pd # 数据处理
import torch
import torch.nn as nn
import torch.optim as optim #  构建和训练LSTM模型
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, mean_absolute_error  # 数据分割和评估指标计算
import matplotlib.pyplot as plt # 绘制图表

# 检查是否有可用的 GPU 设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用的设备: {device}")
print("开始训练...")

# 读取数据
file_path = 'C:\\Users\\严如梦\\Desktop\\文献2023\\康缘\\其他\\300701\\副本T3007A.xlsx'
sheet_name = 'Sheet2'
data = pd.read_excel(file_path, sheet_name=sheet_name)

# 特征提取，提取从第三列到倒数第二列的所有列名作为特征列
feature_columns = data.columns[2:-1]
batch1 = data[data.iloc[:, 1] == 'Z241028-1']
batch2 = data[data.iloc[:, 1] == 'Z241029-1']


# 定义 preprocess_batch 函数将数据转换为适合模型输入的格式
def preprocess_batch(batch):
    X = batch.iloc[:, 2:-1].apply(pd.to_numeric, errors='coerce').values
    Y = batch.iloc[:, -1].apply(pd.to_numeric, errors='coerce').values  #最后一列（即目标变量列）
    return X, Y     #values 将目标变量列的数据转换为NumPy数组

X_batch1, Y_batch1 = preprocess_batch(batch1)
X_batch2, Y_batch2 = preprocess_batch(batch2)

batches = [(X_batch1, Y_batch1), (X_batch2, Y_batch2)]
batch_names = ['Z241028-1', 'Z241029-1']

# 定义 LSTM 神经网络模型
class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_units, num_layers):
        # input_size: 输入特征的维度，即每个时间步长的特征数量。
        # hidden_units: 隐藏层单元的数量，即LSTM层中的神经元数量。
        # num_layers: LSTM层的数量，即堆叠的LSTM层数。
        super(LSTMModel, self).__init__()
        self.num_directions = 2  # 单向LSTM
        self.lstm = nn.LSTM(input_size, hidden_units, num_layers, batch_first=True)
        #  创建一个LSTM层，其参数为输入特征维度、隐藏单元数量、层数，
        #  以及 batch_first=True，表示输入和输出张量的第一个维度是批次大小
        self.fc = nn.Linear(hidden_units, 1)
        # 创建一个全连接层（也称为线性层），它将LSTM的输出转换为最终的预测值。
        # 这里的 hidden_units 是输入特征维度，
        # 而 1 是输出特征维度，因为我们通常在回归问题中预测一个单一的值。

    def forward(self, x):
        out, _ = self.lstm(x)
        # 将输入数据x传递给LSTM层，得到输出out和隐藏状态 _ 。
        # 在这里，我们只关心输出out，而不关心隐藏状态，所以用 _ 来忽略它
        out = self.fc(out[:, -1, :])
        #从LSTM层的输出中取出每个序列的最后一个时间步长的输出（因为 batch_first=True，
        # 所以 out 的形状是 (batch_size, sequence_length, hidden_units)）
        # 并将其传递给全连接层 fc。out[:, -1, :] 表示选择所有批次中最后一个时间步的输出。
        return out
        #  返回模型的最终输出，即全连接层的输出


# 定义目标函数，训练lstm模型
def train_lstm_model(hidden_units, num_layers, X, Y):
    # X: 特征数据，形状为 (samples, time_steps, features)。
    # Y: 目标数据，形状为 (samples, 1)。

    # 数据转换和分割:
    X = torch.tensor(X, dtype=torch.float32).unsqueeze(-1).to(device)
    Y = torch.tensor(Y, dtype=torch.float32).unsqueeze(-1).to(device)
    X_train, X_val, Y_train, Y_val = train_test_split(X, Y, test_size=0.2, random_state=42)


    # 创建模型和优化器,模型初始化:
    model = LSTMModel(input_size=X.shape[1], hidden_units=hidden_units, num_layers=num_layers).to(device)
    # 创建一个 LSTMModel 实例，输入尺寸为 X 的第三个维度（特征数量），隐藏单元数量为 hidden_units，层数为 num_layers
    print(X.shape)

    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    #  训练过程:
    train_losses = []
    val_losses = []

    for epoch in range(300):
        model.train()   # 将模型设置为训练模式
        optimizer.zero_grad()   # 清零梯度，为反向传播准备。
        outputs = model(X_train)    #  将训练数据传递给模型，得到预测输出
        loss = criterion(outputs, Y_train)  # 计算预测输出和真实目标之间的MSE损失
        loss.backward() # 反向传播损失，计算梯度
        optimizer.step()    # 根据梯度更新模型参数。
        train_losses.append(loss.item())

        if epoch % 10 == 0:  # 每10个epoch验证一次
            model.eval()
            with torch.no_grad():
                Y_pred = model(X_val)
                val_loss = mean_squared_error(Y_val.cpu().numpy(), Y_pred.cpu().numpy())
            val_losses.append(val_loss)

    return train_losses, val_losses

# 主程序
results = []
for i, (X, Y) in enumerate(batches):
    train_losses, val_losses = train_lstm_model(32, 2, X, Y)

    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号'-'显示为方框的bug

    plt.figure()
    plt.plot(train_losses, label='训练损失')
    plt.plot(val_losses, label='验证损失')
    plt.legend()
    plt.title(f'批号 {batch_names[i]} 损失曲线')
    plt.show()

    print(X.shape)
    # 训练模型
    model = LSTMModel(input_size=X.shape[1], hidden_units=32, num_layers=2).to(device)

    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(200):
        model.train()
        optimizer.zero_grad()
        outputs = model(torch.tensor(X, dtype=torch.float32).unsqueeze(-1).to(device))
        loss = criterion(outputs, torch.tensor(Y, dtype=torch.float32).unsqueeze(-1).to(device))
        loss.backward()
        optimizer.step()

    model.eval()
    with torch.no_grad():
        Y_pred = model(torch.tensor(X, dtype=torch.float32).unsqueeze(-1).to(device)).flatten()
        mse = mean_squared_error(Y.cpu().numpy(), Y_pred.cpu().numpy())
        mae = mean_absolute_error(Y.cpu().numpy(), Y_pred.cpu().numpy())
        r2 = 1 - mse / np.var(Y.cpu().numpy())
        results.append((mse, mae, r2))

    plt.figure()
    plt.plot(Y.cpu().numpy(), label='实际值', color='red')
    plt.plot(Y_pred.cpu().numpy(), label='预测值', color='blue')
    plt.legend()
    plt.title(f'批号 {batch_names[i]}：实际值 vs 预测值')
    plt.show()

# 打印结果
best_index = np.argmin([r[0] for r in results])
print(f"最佳模型: 批号 {batch_names[best_index]}")

